import platgo as pg
import numpy as np
from typing import Tuple
from scipy.spatial.distance import cdist

"""
Coevolutionary constrained multi-objective optimization framework
type --- 1 --- Type of operator (1. GA 2. DE)

------------------------------- Reference --------------------------------
Y. Tian, T. Zhang, J. Xiao, X. Zhang, and Y. Jin, A coevolutionary
framework for constrained multi-objective optimization problems, IEEE
Transactions on Evolutionary Computation, 2020.
"""


class CCMO(pg.Algorithm):

    type: dict = {'single': False, 'multi': True, 'many': False, 'real': True, 'binary': True, 'permutation': True,
                  "large": False, 'expensive': False, 'constrained': True, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': False}

    def __init__(self, maxgen: int, problem: pg.Problem, operator: int = 1) -> None:
        super().__init__(maxgen=maxgen, problem=problem)
        self.name = "CCMO"
        self.ope = operator
        if operator == 1:
            self.xov = pg.operators.XovSbx(half=True)
        elif operator == 2:
            self.xov = pg.operators.DE(F=0.5)
        self.mut = pg.operators.MutPol(problem)

    def go(self, N: int = None, pop: pg.Population = None):
        assert N or pop, "N and population can't be both None"
        if pop is None:
            pop1 = self.problem.init_pop(N)
        else:
            pop1 = pop
            self.problem.N = pop.decs.shape[0]

        self.problem.cal_obj(pop1)
        self.problem.cal_cv(pop1)
        fitness1 = cal_fitness(pop1.objv, pop1.cv)
        pop2 = self.problem.init_pop()
        self.problem.cal_obj(pop2)
        fitness2 = cal_fitness(pop2.objv)

        while self.not_terminal(pop1):

            if self.ope == 1:
                matingpool = pg.utils.tournament_selection(2, pop1.N, fitness1)
                p1, p2 = pg.utils.random_selection(pop1[matingpool])
                offspring1 = self.xov(pop1, p1, p2)
                offspring1 = self.mut(offspring1)
                self.problem.cal_obj(offspring1)
                self.problem.cal_cv(offspring1)
                matingpool = pg.utils.tournament_selection(2, pop2.N, fitness2)
                p1, p2 = pg.utils.random_selection(pop1[matingpool])
                offspring2 = self.xov(pop2, p1, p2)
                offspring2 = self.mut(offspring2)
                self.problem.cal_obj(offspring2)
                self.problem.cal_cv(offspring2)

            elif self.ope == 2:
                matingpool = pg.utils.tournament_selection(2, 2*pop1.N, fitness1)
                offspring1 = self.xov(pop1, pop1[matingpool[:len(matingpool)//2]], pop1[matingpool[len(matingpool)//2:len(matingpool//2*2)]])
                offspring1 = self.mut(offspring1)
                self.problem.cal_obj(offspring1)
                self.problem.cal_cv(offspring1)
                matingpool = pg.utils.tournament_selection(2, 2*pop2.N, fitness2)
                offspring2 = self.xov(pop2, pop2[matingpool[:len(matingpool)//2]], pop2[matingpool[len(matingpool)//2:len(matingpool//2*2)]])
                offspring2 = self.mut(offspring2)
                self.problem.cal_obj(offspring2)
                self.problem.cal_cv(offspring2)

            temp1 = pop1 + offspring1 + offspring2
            pop1, fitness1 = enviromnent_selection(temp1, pop1.N, True)
            temp2 = pop2 + offspring1 + offspring2
            pop2, fitness2 = enviromnent_selection(temp2, pop2.N, False)
        return pop1


def cal_fitness(objv: np.ndarray, cv: np.ndarray = None) -> np.ndarray:
    N = objv.shape[0]
    if cv is None:
        cv = np.zeros((N, 1))
    else:
        cv = np.where(cv > 0, cv, 0)

    # detect the dominance relation between each two solutions
    dominance = np.zeros((N, N), dtype=bool)
    for i in range(N - 1):
        for j in range(i + 1, N):
            if cv[i] < cv[j]:
                dominance[i][j] = True
            elif cv[i] > cv[j]:
                dominance[j][i] = True
            else:
                k = int(np.any(objv[i] < objv[j])) - int(np.any(objv[i] > objv[j]))
                if k == 1:
                    dominance[i][j] = True
                elif k == -1:
                    dominance[j][i] = True

    # 计算S(i)
    s = np.sum(dominance, axis=1)
    # 计算R(i)
    r = np.zeros(N)
    for i in range(N):
        r[i] = np.sum(s[dominance[:, i]])
    # 计算D(i)
    distance = cdist(objv, objv)
    distance[np.eye(len(distance), dtype=bool)] = np.inf
    distance = np.sort(distance, axis=1)
    d = 1 / (distance[:, int(np.floor(np.sqrt(N)) - 1)] + 2)
    # 计算fitness
    fitness = r + d

    return fitness


def enviromnent_selection(pop: pg.Population, N: int, is_origin: bool) -> Tuple[pg.Population, np.ndarray]:
    if is_origin:
        fitness = cal_fitness(pop.objv, pop.cv)
    else:
        fitness = cal_fitness(pop.objv)

    # 环境选择
    next = fitness < 1
    if np.sum(next) < N:
        rank = np.argsort(fitness)
        next[rank[:N]] = True
    elif np.sum(next) > N:
        delete = truncation(pop[next].objv, np.sum(next) - N)
        temp = np.argwhere(next)
        next[temp[delete]] = False

    pop = pop[next]
    fitness = fitness[next]

    ind = np.argsort(fitness)
    pop = pop[ind]
    fitness = fitness[ind]

    return pop, fitness


def truncation(objv: np.ndarray, K: int):
    # 截断策略
    distance = cdist(objv, objv)
    distance[np.eye(len(distance), dtype=bool)] = np.inf
    delete = np.zeros(len(objv), dtype=bool)

    while np.sum(delete) < K:
        remain = np.argwhere(~delete).flatten()
        temp = distance[remain]
        temp = temp[:, remain]
        temp = np.sort(temp, axis=1)
        _, rank = np.unique(temp, return_index=True, axis=0)
        delete[remain[rank[0]]] = True

    return delete
